from pyspark.sql import SparkSession
from pyspark.sql.functions import countDistinct, count, col
from pyspark.sql.types import IntegerType
import matplotlib.pyplot as plt
import pandas as pd
import os # 导入 os 模块

# --- 输出文件夹定义 ---
output_dir = "output" # 定义输出文件夹名称
# ---------------------

# --- Matplotlib 中文显示配置 ---
# 尝试使用 macOS 常用的 'Heiti TC' 字体
# 如果系统没有，可能需要更换为其他已安装的中文字体，如 'PingFang SC', 'Arial Unicode MS' 等
try:
    plt.rcParams['font.sans-serif'] = ['Heiti TC']  # 指定默认字体为 Heiti TC
    plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题
except Exception as e:
    print(f"无法设置中文字体'Heiti TC'，尝试回退到 PingFang SC 或 Arial Unicode MS。图表中的中文可能无法正常显示: {e}")
    # 可以添加更多备选字体尝试
    try:
        plt.rcParams['font.sans-serif'] = ['PingFang SC'] # 尝试 PingFang SC
        plt.rcParams['axes.unicode_minus'] = False
        print("已尝试设置为 PingFang SC")
    except Exception as e2:
        print(f"也无法设置 PingFang SC: {e2}，尝试 Arial Unicode MS...")
        try:
            plt.rcParams['font.sans-serif'] = ['Arial Unicode MS'] # 尝试 Arial Unicode MS
            plt.rcParams['axes.unicode_minus'] = False
            print("已尝试设置为 Arial Unicode MS")
        except Exception as e3:
             print(f"也无法设置 Arial Unicode MS: {e3}")
             print("请确保系统中安装了支持中文的字体（如 Heiti TC, PingFang SC, SimHei, Microsoft YaHei 等），并在代码中正确指定字体名称。")
# -----------------------------

# 1. 初始化 SparkSession
spark = SparkSession.builder \
    .appName("TaobaoUserAnalysisWithPlots") \
    .getOrCreate()

print("SparkSession 已创建")

# 2. 加载数据集
# 注意: 请确保文件路径正确，并且 Spark 有权限访问
# 使用小数据集进行测试
#file_path_small = "file:///Users/guopeiran/Code/VSCode/Homework/small_user_log.csv"
# 完整数据集路径
file_path_full = "file:///Users/guopeiran/Code/VSCode/Homework/user_log_new.csv"

# 加载数据
df = spark.read.csv(file_path_full, header=True, inferSchema=True) # 使用完整数据集
#df = spark.read.csv(file_path_small, header=True, inferSchema=True) # 使用小数据集

print(f"数据集已从 {file_path_full} 加载") # 修改这里，打印完整数据集路径

# 3. 数据预处理/类型转换
# 确保关键列是整数类型，以便进行过滤和聚合
try:
    df = df.withColumn("action", col("action").cast(IntegerType()))
    df = df.withColumn("brand_id", col("brand_id").cast(IntegerType()))
    df = df.withColumn("gender", col("gender").cast(IntegerType()))
    df = df.withColumn("month", col("month").cast(IntegerType())) # 确保month和day也是整数
    df = df.withColumn("day", col("day").cast(IntegerType()))
    # 年龄范围也转为整数，便于后续处理
    df = df.withColumn("age_range", col("age_range").cast(IntegerType()))
    print("关键列类型转换完成")
    # df.printSchema() # 取消注释以查看Schema
    # 清理操作：去除后续分析中无法解释的 gender 和 age_range 值 (NULL, 0, 2 for gender; NULL, 0 for age_range)
    # 注意： brand_id 的 NULL 可能在特定任务过滤时自动处理，但 gender 和 age_range 用于分组，最好先处理
    df_cleaned = df.filter(col("gender").isin([0, 1])) # 只保留 0 (女), 1 (男)
    df_cleaned = df_cleaned.filter(col("age_range").isin([1, 2, 3, 4, 5, 6, 7, 8])) # 只保留定义的年龄段
    print(f"数据清洗后剩余行数: {df_cleaned.count()} (移除了未知性别和年龄范围)")
except Exception as e:
    print(f"类型转换或清洗时发生错误: {e}")
    print("请检查CSV文件的列名和数据格式是否正确。")
    spark.stop()
    exit()

# --- 分析任务 ---

# 用于存储结果的字典
results = {}

# 任务一：计算总数据行数 (使用原始df)
try:
    total_rows = df.count()
    results["total_rows"] = total_rows
    print(f"\n任务一：表内总行数 (原始数据): {total_rows}")
except Exception as e:
    print(f"\n任务一执行失败: {e}")
    results["total_rows"] = "计算失败"

# 在清洗后的数据上进行后续分析 (df_cleaned)
df_double11 = df_cleaned.filter((col("month") == 11) & (col("day") == 11))
print(f"\n筛选双十一当天数据 (已清洗)，剩余行数: {df_double11.count()}")

df_double11_purchases = df_double11.filter(col("action") == 2)
print(f"筛选双十一当天购买行为数据，剩余行数: {df_double11_purchases.count()}")

# 任务二：查询双11那天有多少人购买了商品（用户ID去重）
try:
    distinct_buyers_double11 = df_double11_purchases \
        .agg(countDistinct("user_id").alias("unique_buyers")) \
        .collect()[0]["unique_buyers"]
    results["distinct_buyers_double11"] = distinct_buyers_double11
    print(f"\n任务二：双11购买商品的不重复用户数: {distinct_buyers_double11}")
except Exception as e:
    print(f"\n任务二执行失败: {e}")
    results["distinct_buyers_double11"] = "计算失败"

# 任务三：给定品牌brand_id为2661，求双11当天此品牌商品的销售数量
try:
    brand_id_target = 2661
    brand_sales_double11 = df_double11_purchases.filter(col("brand_id") == brand_id_target) \
        .agg(count("*").alias("brand_sales")) \
        .collect()[0]["brand_sales"]
    results["brand_sales_double11"] = brand_sales_double11
    print(f"\n任务三：双11品牌 {brand_id_target} 的销售数量: {brand_sales_double11}")
except Exception as e:
    print(f"\n任务三执行失败: {e}")
    results["brand_sales_double11"] = "计算失败"

# 任务四：查询双11那天女性购买商品的数量 (gender=0)
female_purchases_double11 = "计算失败" # 初始化
try:
    # 直接从已筛选的双十一购买数据中计算
    female_purchases_double11 = df_double11_purchases.filter(col("gender") == 0) \
        .agg(count("*").alias("female_purchases")) \
        .collect()[0]["female_purchases"]
    results["female_purchases_double11"] = female_purchases_double11
    print(f"\n任务四：双11女性购买商品数量: {female_purchases_double11}")
except Exception as e:
    print(f"\n任务四执行失败: {e}")
    results["female_purchases_double11"] = "计算失败"

# 任务五：根据收货地址，统计各个省双11那天用户下单的总次数，并进行排序
province_orders_double11_df = None # 初始化变量
try:
    province_orders_double11 = df_double11_purchases \
        .groupBy("province") \
        .agg(count("*").alias("order_count")) \
        .orderBy(col("order_count").desc())

    print("\n任务五：双11各省用户下单总次数统计 (降序):")
    province_orders_double11.show() # 显示排序结果

    # --- 图表生成：任务五 ---
    print("\n正在为任务五生成图表...")
    # 收集结果到 Pandas DataFrame 以便绘图
    province_orders_double11_pd = province_orders_double11.limit(15).toPandas() # 只绘制前15个省份

    if not province_orders_double11_pd.empty:
        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)
        save_path = os.path.join(output_dir, "province_orders_double11_top15.png") # 构建保存路径

        plt.figure(figsize=(12, 7))
        plt.bar(province_orders_double11_pd['province'], province_orders_double11_pd['order_count'])
        plt.xlabel("省份")
        plt.ylabel("下单次数")
        plt.title("双十一各省份用户下单次数 Top 15")
        plt.xticks(rotation=45, ha='right') # 旋转x轴标签以防重叠
        plt.tight_layout() # 调整布局防止标签重叠
        plt.savefig(save_path) # 保存图表到指定路径
        print(f"任务五图表已保存为 {save_path}")
        # plt.show() # 如果希望在运行时显示图表，取消此行注释
        plt.close() # 关闭当前图形，释放内存
    else:
        print("任务五无数据可供绘图。")
    # -----------------------

except Exception as e:
    print(f"\n任务五执行或绘图失败: {e}")

# --- 额外图表生成 ---

# 1. 双十一用户行为分布 (饼图)
try:
    print("\n正在生成双十一用户行为分布饼图...")
    action_dist_double11 = df_cleaned.groupBy("action") \
        .agg(count("*").alias("action_count")) \
        .orderBy("action")

    action_dist_pd = action_dist_double11.toPandas()

    if not action_dist_pd.empty:
        action_labels = {0: '点击', 1: '加购', 2: '购买', 3: '关注'}
        action_dist_pd['action_label'] = action_dist_pd['action'].map(action_labels)

        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)
        save_path = os.path.join(output_dir, "action_distribution_double11.png")

        plt.figure(figsize=(8, 8))
        plt.pie(action_dist_pd['action_count'], labels=action_dist_pd['action_label'], autopct='%1.1f%%', startangle=140)
        plt.title('双十一用户行为分布')
        plt.tight_layout()
        plt.savefig(save_path)
        print(f"用户行为分布饼图已保存为 {save_path}")
        # plt.show()
        plt.close()
    else:
        print("无用户行为数据可供绘图。")

except Exception as e:
    print(f"\n生成用户行为分布图失败: {e}")

# 2. 双十一购买用户性别分布 (条形图) - 增强版
try:
    print("\n正在生成双十一购买用户性别分布条形图 (带数值标注)...")
    gender_purchase_double11 = df_double11_purchases \
        .groupBy("gender") \
        .agg(count("*").alias("purchase_count"))

    gender_purchase_pd = gender_purchase_double11.toPandas()

    if not gender_purchase_pd.empty:
        gender_labels = {0: '女性', 1: '男性'}
        gender_purchase_pd['gender_label'] = gender_purchase_pd['gender'].map(gender_labels)
        # 确保包含女性和男性，即使某一方计数为0
        if '女性' not in gender_purchase_pd['gender_label'].values:
            gender_purchase_pd = pd.concat([gender_purchase_pd, pd.DataFrame([{'gender_label': '女性', 'purchase_count': 0}])], ignore_index=True)
        if '男性' not in gender_purchase_pd['gender_label'].values:
             gender_purchase_pd = pd.concat([gender_purchase_pd, pd.DataFrame([{'gender_label': '男性', 'purchase_count': 0}])], ignore_index=True)
        gender_purchase_pd = gender_purchase_pd.sort_values(by='gender_label') # 排序以便颜色对应

        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)
        save_path = os.path.join(output_dir, "gender_purchase_distribution_double11.png")

        plt.figure(figsize=(7, 6)) # 稍微调整尺寸
        bars = plt.bar(gender_purchase_pd['gender_label'], gender_purchase_pd['purchase_count'], color=['pink', 'lightblue'])

        # 添加数值标签
        for bar in bars:
            yval = bar.get_height()
            plt.text(bar.get_x() + bar.get_width()/2.0, yval, int(yval), va='bottom', ha='center') # va='bottom' 让文字在条形上方

        plt.xlabel("性别")
        plt.ylabel("购买次数")
        # 标题中加入总购买次数信息 (如果任务二成功)
        total_purchases_str = f" (总购买次数: {df_double11_purchases.count()})" if df_double11_purchases else ""
        plt.title(f"双十一购买用户性别分布{total_purchases_str}")
        plt.tight_layout()
        plt.savefig(save_path)
        print(f"购买用户性别分布图 (带数值) 已保存为 {save_path}")
        # plt.show()
        plt.close()
    else:
        print("无购买用户性别数据可供绘图。")

except Exception as e:
    print(f"\n生成购买用户性别分布图失败: {e}")

# 3. 双十一购买用户年龄段分布 (条形图)
try:
    print("\n正在生成双十一购买用户年龄段分布条形图...")
    age_purchase_double11 = df_cleaned.filter(col("action") == 2) \
        .groupBy("age_range") \
        .agg(count("*").alias("purchase_count")) \
        .orderBy("age_range") # 按年龄段排序

    age_purchase_pd = age_purchase_double11.toPandas()

    if not age_purchase_pd.empty:
        age_labels = {
            1: '<18', 2: '[18,24]', 3: '[25,29]', 4: '[30,34]',
            5: '[35,39]', 6: '[40,49]', 7: '>=50', 8: '>=50' # 合并7和8
        }
        # 合并 age_range 7 和 8 的计数
        if 7 in age_purchase_pd['age_range'].values and 8 in age_purchase_pd['age_range'].values:
             count_7 = age_purchase_pd.loc[age_purchase_pd['age_range'] == 7, 'purchase_count'].iloc[0]
             count_8 = age_purchase_pd.loc[age_purchase_pd['age_range'] == 8, 'purchase_count'].iloc[0]
             age_purchase_pd.loc[age_purchase_pd['age_range'] == 7, 'purchase_count'] = count_7 + count_8
             age_purchase_pd = age_purchase_pd[age_purchase_pd['age_range'] != 8] # 移除第8行
        elif 8 in age_purchase_pd['age_range'].values: # 如果只有8，把它映射到7的标签
             age_purchase_pd.loc[age_purchase_pd['age_range'] == 8, 'age_range'] = 7


        age_purchase_pd['age_label'] = age_purchase_pd['age_range'].map(age_labels)
        # 确保按正确的年龄顺序绘图
        age_order = ['<18', '[18,24]', '[25,29]', '[30,34]', '[35,39]', '[40,49]', '>=50']
        age_purchase_pd['age_label'] = pd.Categorical(age_purchase_pd['age_label'], categories=age_order, ordered=True)
        age_purchase_pd = age_purchase_pd.sort_values('age_label')

        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)
        save_path = os.path.join(output_dir, "age_purchase_distribution_double11.png")

        plt.figure(figsize=(10, 6))
        plt.bar(age_purchase_pd['age_label'], age_purchase_pd['purchase_count'])
        plt.xlabel("年龄段")
        plt.ylabel("购买次数")
        plt.title("双十一购买用户年龄段分布")
        plt.tight_layout()
        plt.savefig(save_path)
        print(f"购买用户年龄段分布图已保存为 {save_path}")
        # plt.show()
        plt.close()
    else:
        print("无购买用户年龄段数据可供绘图。")

except Exception as e:
    print(f"\n生成购买用户年龄段分布图失败: {e}")

# --- 生成结果总结文件 ---
try:
    summary_file_path = os.path.join(output_dir, "summary_results.txt")
    os.makedirs(output_dir, exist_ok=True) # 确保目录存在
    with open(summary_file_path, 'w', encoding='utf-8') as f:
        f.write("淘宝用户行为分析结果摘要 (基于 Taobao.md 要求)\n")
        f.write("===============================================\n\n")
        f.write(f"任务一：数据集总行数 (原始): {results.get('total_rows', '未计算')}\n")
        f.write(f"任务二：双十一当天购买商品的不重复用户数: {results.get('distinct_buyers_double11', '未计算')}\n")
        f.write(f"任务三：双十一当天品牌 {brand_id_target} 的销售数量: {results.get('brand_sales_double11', '未计算')}\n")
        f.write(f"任务四：双十一当天女性购买商品数量: {results.get('female_purchases_double11', '未计算')}\n")
        f.write("\n任务五：双十一当天各省用户下单总次数统计: 请查看 output/province_orders_double11_top15.png 图表及控制台输出\n")
        f.write("\n其他分析图表请查看 output/ 文件夹下的其他 .png 文件。\n")
    print(f"\n结果摘要已写入文件: {summary_file_path}")
except Exception as e:
    print(f"\n写入结果摘要文件失败: {e}")
# ----------------------

# 6. 停止 SparkSession
spark.stop()
print("\nSparkSession 已停止") 